# management/commands/train_model.py
from django.core.management.base import BaseCommand
from django.utils import timezone


class Command(BaseCommand):
    help = '训练成绩预测模型'

    def add_arguments(self, parser):
        parser.add_argument(
            '--data-source',
            type=str,
            choices=['simulated', 'database'],
            default='simulated',
            help='数据源: simulated(模拟数据) 或 database(数据库数据)'
        )
        parser.add_argument(
            '--model-path',
            type=str,
            help='模型保存路径（默认: ml/models/score_predictor.pkl）'
        )
        parser.add_argument(
            '--sample-size',
            type=int,
            default=1000,
            help='模拟数据样本数量（仅simulated模式有效）'
        )

    def handle(self, *args, **options):
        self.stdout.write('=' * 60)
        self.stdout.write('🚀 开始训练成绩预测模型')
        self.stdout.write('=' * 60)

        try:
            # 导入训练器
            from prediction.ml.trainers.model_trainer import ModelTrainer

            # 创建训练器
            trainer = ModelTrainer()

            # 显示训练信息
            model_info = trainer.get_model_info()
            self.stdout.write(
                self.style.SUCCESS(f'📊 模型状态: {model_info["status"]}')
            )
            self.stdout.write(
                self.style.SUCCESS(f'🔧 特征数量: {model_info["feature_count"]}')
            )

            # 开始训练
            self.stdout.write('🎯 开始模型训练...')
            start_time = timezone.now()

            result = trainer.train()

            training_time = timezone.now() - start_time

            # 显示训练结果
            self.stdout.write('\n' + '=' * 60)
            self.stdout.write('🏆 训练结果')
            self.stdout.write('=' * 60)

            if result['success']:
                self.stdout.write(
                    self.style.SUCCESS('✅ 训练成功!')
                )
                self.stdout.write(
                    f'📈 训练集 R²: {result["train_score"]:.4f}'
                )
                self.stdout.write(
                    f'📈 测试集 R²: {result["test_score"]:.4f}'
                )
                self.stdout.write(
                    f'📏 平均绝对误差: {result["mae"]:.2f} 分'
                )
                self.stdout.write(
                    f'📊 训练样本: {result["training_samples"]} 条'
                )
                self.stdout.write(
                    f'⏱️ 训练耗时: {training_time.total_seconds():.2f} 秒'
                )

                # 显示重要特征
                feature_importance = result.get('feature_importance', {})
                if feature_importance:
                    top_features = sorted(
                        feature_importance.items(),
                        key=lambda x: x[1],
                        reverse=True
                    )[:5]

                    self.stdout.write('\n📋 最重要的5个特征:')
                    for feature, importance in top_features:
                        self.stdout.write(
                            f'   {feature}: {importance:.3f}'
                        )

            else:
                self.stdout.write(
                    self.style.ERROR(f'❌ 训练失败: {result.get("error", "未知错误")}')
                )

        except ImportError as e:
            self.stdout.write(
                self.style.ERROR(f'❌ 导入失败: {e}')
            )
            self.stdout.write('💡 请检查:')
            self.stdout.write('   1. prediction.ml 目录结构是否正确')
            self.stdout.write('   2. 所有 __init__.py 文件是否存在')
            self.stdout.write('   3. 依赖包是否安装 (scikit-learn, joblib等)')

        except Exception as e:
            self.stdout.write(
                self.style.ERROR(f'❌ 训练过程出错: {e}')
            )
            import traceback
            self.stdout.write(traceback.format_exc())